#
# This file is a part of the open-eBackup project.
# This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0.
# If a copy of the MPL was not distributed with this file, You can obtain one at
# http://mozilla.org/MPL/2.0/.
#
# Copyright (c) [2024] Huawei Technologies Co.,Ltd.
#
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
#

import json
import os
import pipes
import re
import stat
import tempfile
import uuid
from base64 import b64decode
import pexpect
import psutil
from clickhouse import log
from common.common import check_command_injection, retry_when_exception
from common.const import AuthType
from common.file_common import change_path_permission
from common.util.check_utils import is_port, is_ip_address
from clickhouse.common.clickhouse_constants import PluginPathConstant, KerberosConstant, \
    CommandReturnCode, AccessPermission, KerberosCmdRetKey, SysData
from clickhouse.schemas.clickhouse_schemas import Database, Disk, Cluster, Table, Replicas
from common.util.check_user_utils import check_path_owner
from common.util.cmd_utils import cmd_format

declname_re = re.compile(r'^[\w.]+$')
path_re = re.compile(r"^[^|;&$`<>!\n\\]+$")
user_re = re.compile(r"^(?!-)[\w\s-]{3,64}$")


class ClickHouseClient(object):
    """
    功能描述：clickhouse客户端
    """

    def __init__(self, pid: str, app_extend_info, env_prefix: str, auth_extend_info):
        """
        功能描述：__init__ 构造函数
        参数：
        @pid：请求ID
        @app_extend_info：应用连接信息
        @env_prefix：敏感信息，环境变量前缀
        @auth_extend_info：鉴权auth信息
        返回值：无
        """
        log.info(f'client init begin, pid: {pid}, env_prefix: {env_prefix}')
        self.pid = pid
        self.env_prefix = env_prefix
        self.process_name = "clickhouse"
        self.process_owner = None
        self.process_exe = None
        self.client_path = os.path.realpath(app_extend_info['clientPath'])
        if not path_re.match(self.client_path):
            log.error(f'the client_path is Illegal {self.client_path}')
            raise Exception(f'the client_path is Illegal {self.client_path}')
        self.listen_ip = app_extend_info['ip']
        if not is_ip_address(self.listen_ip):
            log.error(f'the ip_address is Illegal {self.listen_ip}')
            raise Exception(f'the ip_address is Illegal {self.listen_ip}')
        self.listen_port = app_extend_info['port']
        if not is_port(self.listen_port):
            log.error(f'the port is Illegal {self.listen_port}')
            raise Exception(f'the port is Illegal {self.listen_port}')
        self.authenticated = False
        self.db_user = self.get_env_variable(f'{self.env_prefix}_authKey_{self.pid}')
        self.check_db_user(self.db_user)
        self.db_pwd = self.get_env_variable(f'{self.env_prefix}_authPwd_{self.pid}')
        if check_command_injection(self.db_user):
            log.error('the database login information is Illegal.')
            raise Exception('the database login information is Illegal.')
        self.extend_info_principal = self.get_env_variable(f'{self.env_prefix}_extendInfo_principal_{self.pid}')
        self.extend_info_keytab = self.get_env_variable(f'{self.env_prefix}_extendInfo_keytab_{self.pid}')
        self.extend_info_secret = self.get_env_variable(f'{self.env_prefix}_extendInfo_secret_{self.pid}')
        # 鉴权类型: NO_AUTO:0,APP_PASSWORD:2,KERBEROS:5
        self.auth_type = int(self.get_env_variable(f'{self.env_prefix}_authType_{self.pid}'))
        if self.auth_type is AuthType.KERBEROS.value:
            # Kerberos配置方式: 密码：password_model,Keytab文件: keytab_model
            self.krb_config_model = auth_extend_info["configModel"]
            self.krb_krb5_config_value = auth_extend_info["krb5Conf"]
            self.krb_keytab_file = PluginPathConstant.TMP_PATH.value + f'user_{self.pid}.keytab'
            self.krb_krb5_file = PluginPathConstant.TMP_PATH.value + f'krb5_{self.pid}.conf'
        self.init_process_params()
        self.init_kerberos_files()
        log.info(f'client init end, pid: {self.pid}')

    def init_process_params(self):
        """
        功能描述：初始化clickhouse进程所属用户及进程绝对路径
        参数：无
        返回值：无
        """
        processes = psutil.process_iter()
        for process in processes:
            try:
                if self.process_name == process.name():
                    self.process_owner = process.username()
                    self.process_exe = process.exe()
                    break
            except Exception as e:
                log.error(e, exc_info=True)
        if self.process_owner is None or self.process_exe is None:
            raise Exception(f'init_process_params failed')

    def init_kerberos_files(self):
        """
        功能描述：创建krb5_{self.pid}.conf user_{self.pid}.keytab临时文件
        参数：无
        返回值：无
        """
        if self.auth_type is AuthType.KERBEROS.value:
            self.write_file(self.krb_krb5_file, b64decode(self.krb_krb5_config_value), 'wb')
            if self.krb_config_model == KerberosConstant.KEYTAB_MODEL.value:
                self.write_file(self.krb_keytab_file, b64decode(self.extend_info_keytab), 'wb')

    def remove_kerberos_files(self):
        """
        功能描述：移除krb5_{self.pid}.conf user_{self.pid}.keytab临时文件
        参数：无
        返回值：无
        """
        if self.auth_type is AuthType.KERBEROS.value:
            if os.path.exists(self.krb_krb5_file):
                os.remove(self.krb_krb5_file)
            if os.path.exists(self.krb_keytab_file):
                os.remove(self.krb_keytab_file)

    def execute_kerberos_auth_cmd(self):
        """
        功能描述：执行kerberos鉴权命令
        参数：无
        返回值：无
        """
        child = None
        try:
            # 校验环境变量文件是否存在
            if not os.path.exists(os.path.join(self.client_path, "bigdata_env")):
                log.error(f"{self.client_path}/bigdata_env not exist")
                raise Exception(f"bigdata_env not exist")

            # 校验环境变量文件用户属主
            if not check_path_owner(os.path.join(self.client_path, "bigdata_env"), ["root", f"{self.process_owner}"]):
                log.error(
                    f"{self.client_path}/bigdata_env check_path_owner failed, expect root or {self.process_owner}")
                raise Exception(
                    f"{self.client_path}/bigdata_env check_path_owner failed, expect root or {self.process_owner}")

            if self.krb_config_model == KerberosConstant.PASSWORD_MODEL.value:
                krb_auth_cmd = cmd_format("su - {} -c 'export KRB5_CONF={} && source {}/bigdata_env && kinit {}'",
                                          self.process_owner, self.krb_krb5_file, self.client_path, self.db_user)
                child = pexpect.spawn(krb_auth_cmd, encoding='utf-8')
                child.expect(['Password for', pexpect.EOF])
                if KerberosCmdRetKey.CLIENT_NOT_FOUND.value in child.before:
                    log.error(f'error info:{child.before}')
                    raise Exception(KerberosConstant.KINIT_FAILED.value)
                child.sendline(self.db_pwd)
                child.expect(pexpect.EOF)
                if KerberosCmdRetKey.INFORMATION_ERROR.value in child.before:
                    log.error(f'error info:{child.before}')
                    raise Exception(KerberosConstant.KINIT_FAILED.value)
                child.sendline('klist')
                child.expect(pexpect.EOF)
            else:
                krb_auth_cmd = cmd_format(
                    "su - {} -c 'export KRB5_CONF={} && source {}/bigdata_env && kinit -kt {} {} && klist'",
                    self.process_owner, self.krb_krb5_file, self.client_path, self.krb_keytab_file, self.db_user)
                child = pexpect.spawn(krb_auth_cmd, encoding='utf-8')
                child.expect(pexpect.EOF)
            password_expired = KerberosCmdRetKey.PASSWORD.value in child.before and \
                               KerberosCmdRetKey.EXPIRED.value in child.before
            if password_expired:
                log.error(f'error info:{child.before}')
                raise Exception(KerberosConstant.KINIT_FAILED.value)
            error_list = [
                KerberosCmdRetKey.NO_CREDENTIALS_CACHE.value, KerberosCmdRetKey.CLIENT_NOT_FOUND.value,
                KerberosCmdRetKey.PREAUTHENTICATION_FAILED.value
            ]
            for error in error_list:
                if error in child.before:
                    log.error(f'error info:{child.before}')
                    raise Exception(KerberosConstant.KINIT_FAILED.value)
        except Exception as e:
            log.error(e, exc_info=True)
            raise e
        finally:
            if child:
                child.close()

    def write_file(self, file_path, content, mode):
        """
        功能描述：写鉴权信息到临时文件，切权限供clickhouse进程用户访问
        参数：无
        返回值：无
        """
        with open(file_path, mode) as file:
            file.write(content)
        change_path_permission(file_path, mode=AccessPermission.RWX_RX_RX.value)

    @retry_when_exception(retry_times=3, delay=5)
    def execute_clickhouse_cmd(self, exec_sql: str):
        """
        功能描述：执行clickhouse命令
        参数：
        @exec_sql：待执行的sql语句
        返回值：
        @code：错误码
        @out：正常流
        @err：异常流
        """

        # 校验环境变量文件是否存在
        if not os.path.exists(os.path.join(self.client_path, "bigdata_env")):
            log.error(f"{self.client_path}/bigdata_env not exist")
            raise Exception(f"bigdata_env not exist")

        # 校验环境变量文件用户属主
        if not check_path_owner(os.path.join(self.client_path, "bigdata_env"), ["root", f"{self.process_owner}"]):
            log.error(
                f"{self.client_path}/bigdata_env check_path_owner failed, expect root or {self.process_owner}")
            raise Exception(
                f"{self.client_path}/bigdata_env check_path_owner failed, expect root or {self.process_owner}")

        exec_sql = pipes.quote(exec_sql)
        exec_sql = exec_sql[1: len(exec_sql) - 1]

        # exec_sql在上面已经转义过了， 无需再格式化。ip和port和dbuser在初始化时已经过检验。process_exe非输入参数，无需校验
        clickhouse_cmd = f"source {pipes.quote(self.client_path)}/bigdata_env && stty -echo " \
                         f"&& {self.process_exe} client " \
                         f"--host \"{self.listen_ip}\" --port \"{self.listen_port}\" " \
                         f"--user \"{self.db_user}\" --password --query=\"{exec_sql}\" --secure"
        if self.auth_type == AuthType.KERBEROS.value:
            if not self.authenticated:
                self.execute_kerberos_auth_cmd()
                self.authenticated = True
        fd, file_name = tempfile.mkstemp(suffix=f'{str(uuid.uuid4())}', prefix=f'{self.pid}')
        file_output = os.fdopen(os.open(file_name, os.O_WRONLY, stat.S_IWUSR | stat.S_IRUSR), 'w')
        command = f"su - {self.process_owner} -c '{clickhouse_cmd}'"
        child = pexpect.spawn(command, encoding='utf-8')
        child.expect(['Password', pexpect.EOF])
        child.sendline(f'{self.db_pwd}')
        child.logfile = file_output
        child.expect(pexpect.EOF)
        child.close()
        file_output.close()
        file_input = open(file_name, 'r')
        out = file_input.read()
        out = out.strip()
        if out.startswith('ClickHouse start in security mod!'):
            out = out.replace('ClickHouse start in security mod!', '', 1)
        out = out.strip()
        file_input.close()
        os.close(fd)
        os.remove(file_name)
        if child.exitstatus == CommandReturnCode.SUCCESS.value:
            return CommandReturnCode.SUCCESS.value, out, ''
        else:
            if KerberosConstant.DOES_NOT_EXIST.value in out:
                raise Exception(KerberosConstant.TABLE_DOES_NOT_EXIST.value)
            log.error(f'execute_clickhouse_cmd failed, exit status: {child.exitstatus}, err: {out}')
            raise Exception(f'execute_clickhouse_cmd failed')

    def query_system_disks(self, disk_name=None):
        """
        功能描述：查询系统磁盘
        参数：
        @disk_name：磁盘名称
        返回值：
        @disks：磁盘列表
        """
        disks = []
        sql = "select name, path from system.disks"
        if disk_name is not None:
            sql = sql + f" where name = '{disk_name}'"
        code, out, err = self.execute_clickhouse_cmd(sql)
        if len(out) > 0:
            for row in out.splitlines():
                disk = Disk()
                columns = row.split()
                disk.name = columns[0]
                disk.path = columns[1]
                disks.append(disk)
        return disks

    def query_system_clusters(self):
        """
        功能描述：查询系统集群
        参数：无
        返回值：
        @clusters：集群列表
        """
        clusters = []
        sql = "select cluster, shard_num, shard_weight, replica_num, host_name, host_address, port, is_local, user " \
              "from system.clusters"
        code, out, err = self.execute_clickhouse_cmd(sql)
        if len(out) > 0:
            for row in out.splitlines():
                columns = row.split()
                cluster = Cluster()
                cluster.cluster = columns[0]
                cluster.shard_num = columns[1]
                cluster.shard_weight = columns[2]
                cluster.replica_num = columns[3]
                cluster.host_name = columns[4]
                cluster.host_address = columns[5]
                cluster.port = columns[6]
                cluster.is_local = columns[7]
                cluster.user = columns[8]
                clusters.append(cluster)
        return clusters

    def query_system_replicas(self, db_name: str, table_name=None, condition=None):
        """
        功能描述：查询复制表副本
        参数：
        @db_name：数据库名称
        @table_name：数据表名称
        返回值：
        @replicas：副本列表
        """
        if not declname_re.match(db_name):
            raise Exception(f'the db_name is Illegal {db_name}')
        replicas = []
        sql = f"select table, engine, zookeeper_path, replica_name, database from system.replicas " \
              f"where database = '{db_name}'"
        if table_name is not None:
            if not declname_re.match(table_name):
                raise Exception(f'the table_name is Illegal {table_name}')
            sql = sql + f" and table = '{table_name}'"
        if condition is not None:
            page_no = int(condition['pageNo'])
            page_size = int(condition['pageSize'])
            if page_size != 0:
                sql = sql + f" ORDER BY table limit {(page_no - 1) * page_size},{page_size}"
        code, out, err = self.execute_clickhouse_cmd(sql)
        if len(out) > 0:
            for row in out.splitlines():
                replica = Replicas()
                columns = row.split()
                replica.table = columns[0]
                replica.engine = columns[1]
                replica.zookeeper_path = columns[2]
                replica.replica_name = columns[3]
                replica.database = columns[4]
                replicas.append(replica)
        return replicas

    def query_system_version(self):
        """
        功能描述：查询系统版本号
        参数：无
        返回值：系统版本号
        """
        code, out, err = self.execute_clickhouse_cmd("select version()")
        return out.splitlines()[0]

    def query_system_settings(self):
        """
        功能描述：查询系统设置，创建表是否可以指定表UUID
        参数：无
        返回值：数值 > 0支持指定UUID
        """
        sql = f"select count() from system.settings WHERE name = 'show_table_uuid_in_table_create_query_if_not_nil'"
        code, out, err = self.execute_clickhouse_cmd(sql)
        return out.splitlines()[0]

    def query_allow_backup_databases(self, db_name=None):
        """
        功能描述：查询可备份的库，部分库和库引擎不支持备份
        参数：无
        @db_name： 数据库名称
        返回值：
        @databases: 库列表
        """
        databases = []
        sql = "select uuid, name, engine from system.databases WHERE name NOT IN ('system', 'INFORMATION_SCHEMA', " \
              "'information_schema') AND engine NOT IN ('MySQL','PostgreSQL')"
        if db_name is not None:
            if not declname_re.match(db_name):
                raise Exception(f'the db_name is Illegal {db_name}')
            sql = sql + f" AND name = '{db_name}'"
        code, out, err = self.execute_clickhouse_cmd(sql)
        if len(out) > 0:
            for row in out.splitlines():
                database = Database()
                columns = row.split()
                database.uuid = columns[0]
                database.name = columns[1]
                database.engine = columns[2]
                databases.append(database)
        return databases

    def query_allow_backup_tables_count(self, db_name):
        """
        功能描述：查询库下可备份的表，临时表除外
        参数：无
        @db_name： 数据库名称
        @table_name： 表名称
        返回值：
        @tables: 表列表
        """
        count = 0
        sql = f"select count(*) from system.tables " \
              f"where is_temporary = 0"
        if not declname_re.match(db_name):
            raise Exception(f'the db_name is Illegal {db_name}')
        sql = sql + f" and database = '{db_name}'"
        code, out, err = self.execute_clickhouse_cmd(sql)
        if len(out) > 0:
            for row in out.splitlines():
                count = int(row)
        return count

    def query_allow_backup_tables(self, db_name=None, table_name=None, condition=None):
        """
        功能描述：查询库下可备份的表，临时表除外
        参数：无
        @db_name： 数据库名称
        @table_name： 表名称
        返回值：
        @tables: 表列表
        """
        tables = []
        sql = f"select uuid, name, engine, database, data_paths, create_table_query from system.tables " \
              f"where is_temporary = 0"
        if db_name is not None:
            if not declname_re.match(db_name):
                raise Exception(f'the db_name is Illegal {db_name}')
            sql = sql + f" and database = '{db_name}'"
        if table_name is not None:
            if not declname_re.match(table_name):
                raise Exception(f'the table_name is Illegal {table_name}')
            sql = sql + f" and name = '{table_name}'"
        if condition is not None:
            page_no = int(condition['pageNo'])
            page_size = int(condition['pageSize'])
            if page_size != 0:
                sql = sql + f" ORDER BY name limit {(page_no - 1) * page_size},{page_size}"
        if self.query_system_settings() != "0":
            sql = sql + f" SETTINGS show_table_uuid_in_table_create_query_if_not_nil=1"
        code, out, err = self.execute_clickhouse_cmd(sql)
        if len(out) > 0:
            for row in out.splitlines():
                table = Table()
                columns = row.split()
                table.uuid = columns[0]
                table.name = columns[1]
                table.engine = columns[2]
                table.database = columns[3]
                table.data_paths = columns[4]
                table.create_table_query = ' '.join(columns[5:len(columns)]).replace("\\'", "'")
                tables.append(table)
        return tables

    def query_tables_by_uuid(self, table_uuid=None):
        """
        功能描述：根据uuid查询库下的表
        参数：无
        @table_uuid： uuid
        返回值：
        @tables: 表列表
        """
        tables = []
        sql = f"select uuid, name, database from system.tables " \
              f"where is_temporary = 0"
        if table_uuid is not None:
            sql = sql + f" and uuid = '{table_uuid}'"
        if self.query_system_settings() != "0":
            sql = sql + f" SETTINGS show_table_uuid_in_table_create_query_if_not_nil=1"
        code, out, err = self.execute_clickhouse_cmd(sql)
        if len(out) > 0:
            for row in out.splitlines():
                table = Table()
                columns = row.split()
                table.uuid = columns[0]
                table.name = columns[1]
                table.database = columns[2]
                tables.append(table)
        return tables

    def freeze_table(self, cluster_name: str, db_name: str, t_name: str, t_engine: str, bak_dir_prefix: str):
        """
        功能描述：冻结表, 只有MergeTree表引擎家族才支持冻结, 中划线冻结生成目录会转义
        参数：
        @cluster_name：集群名
        @db_name：数据库名
        @t_name：表名称
        @t_engine：表引擎
        @bak_dir_name：备份目录名前缀
        返回值：无
        """
        freeze_name = bak_dir_prefix.replace('-', '_') + '_' + db_name + '_' + t_name
        if t_engine.endswith("MergeTree"):
            log.info(f"freeze table, t_name: {t_name}, t_engine: {t_engine}, freeze_name: {freeze_name}")
            if not declname_re.match(t_name):
                raise Exception(f'the db_name is Illegal {t_name}')
            if not declname_re.match(db_name):
                raise Exception(f'the db_name is Illegal {db_name}')
            sql = f"alter table {db_name}.{t_name} on cluster '{cluster_name}' freeze with name '{freeze_name}'"
            try:
                self.execute_clickhouse_cmd(sql)
            except Exception as e:
                log.error(e, exc_info=True)
                return False
        else:
            log.info(f"not support freeze table, t_name: {t_name}, t_engine: {t_engine}")
        return True

    def attach_table_part(self, db_name: str, table_name: str, part_name: str):
        """
        功能描述：挂载表分区
        参数：
        @db_name： 库名称
        @table_name： 表名称
        @part_name： 分区名
        返回值：无
        """
        if not declname_re.match(db_name):
            raise Exception(f'the db_name is Illegal {db_name}')
        if not declname_re.match(table_name):
            raise Exception(f'the db_name is Illegal {table_name}')
        return self.execute_clickhouse_cmd(f"alter table {db_name}.{table_name} attach part '{part_name}'")

    def truncate_table(self, db_name: str, table_name: str, t_engine: str):
        """
        功能描述：清除表数据, 应用目前只支持MergeTree引擎备份, 其他类型不清除
        参数：
        @db_name： 库名称
        @table_name： 表名称
        @t_engine： 表引擎
        返回值：无
        """
        if t_engine.endswith("MergeTree"):
            log.info(f"truncate table, t_name: {table_name}, t_engine: {t_engine}")
            if not declname_re.match(db_name):
                raise Exception(f'the db_name is Illegal {db_name}')
            if not declname_re.match(table_name):
                raise Exception(f'the db_name is Illegal {table_name}')
            self.execute_clickhouse_cmd(f"truncate table {db_name}.{table_name}")
        else:
            log.info(f"not support truncate table, t_name: {table_name}, t_engine: {t_engine}")

    def drop_table(self, db_name: str, table_name: str):
        """
        功能描述：删除数据表, no delay表示立即删除, 不用等后台调度删除
        参数：
        @db_name： 库名称
        @table_name： 表名称
        返回值：无
        """
        if not declname_re.match(db_name):
            raise Exception(f'the db_name is Illegal {db_name}')
        if not declname_re.match(table_name):
            raise Exception(f'the db_name is Illegal {table_name}')
        return self.execute_clickhouse_cmd(f"drop table {db_name}.{table_name} no delay")

    def create_table(self, exec_sql: str):
        """
        功能描述：创建数据表
        参数：
        @exec_sql： 创库表SQL
        返回值：无
        """
        return self.execute_clickhouse_cmd(exec_sql)

    def create_database(self, exec_sql: str):
        """
        功能描述：创建数据库
        参数：
        @exec_sql： 创库SQL
        返回值：无
        """
        return self.execute_clickhouse_cmd(exec_sql)

    def drop_database(self, db_name: str):
        if not declname_re.match(db_name):
            raise Exception(f'the db_name is Illegal {db_name}')
        return self.execute_clickhouse_cmd(f"drop database if exists {db_name} no delay")

    def truncate_table_no_engine(self, db_name: str, table_name: str):
        if not declname_re.match(db_name):
            raise Exception(f'the db_name is Illegal {db_name}')
        if not declname_re.match(table_name):
            raise Exception(f'the db_name is Illegal {table_name}')
        return self.execute_clickhouse_cmd(f"truncate table {db_name}.{table_name}")

    def get_env_variable(self, str_env_variable: str):
        env_variable = ''
        input_str = json.loads(SysData.SYS_STDIN)
        if input_str.get(str_env_variable):
            env_variable = input_str.get(str_env_variable)
        return env_variable

    def check_db_user(self, db_user: str):
        if not user_re.match(db_user):
            raise Exception(f'the db_user is Illegal {db_user}')
